{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-07-19T18:01:57.555057Z",
     "start_time": "2024-07-19T18:01:56.787958Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from openai import OpenAI\n",
    "import json\n",
    "from tqdm.notebook import tqdm\n",
    "import time\n"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "client = OpenAI(api_key=\"your-api-key\")\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-07-19T18:01:57.631011Z",
     "start_time": "2024-07-19T18:01:57.555134Z"
    }
   },
   "id": "1e970ed2a0baa46",
   "execution_count": 2
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "74"
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = []\n",
    "with open(\"./data.json\", \"r\") as f:\n",
    "    for line in f:\n",
    "        data.append(eval(line))\n",
    "        \n",
    "len(data)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-07-19T18:02:36.093888Z",
     "start_time": "2024-07-19T18:02:36.053872Z"
    }
   },
   "id": "84e9bb7a73b6963d",
   "execution_count": 4
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "MODEL_LIMITS = {\n",
    "    \"gpt-3.5-turbo-0125\": 16_385,\n",
    "    \"gpt-4-turbo-2024-04-09\": 128_000,\n",
    "    \"gpt-4o-2024-05-13\": 128_000\n",
    "}\n",
    "\n",
    "# The cost per token for each model input.\n",
    "MODEL_COST_PER_INPUT = {\n",
    "    \"gpt-3.5-turbo-0125\": 0.0000005,\n",
    "    \"gpt-4-turbo-2024-04-09\": 0.00001,\n",
    "    \"gpt-4o-2024-05-13\": 0.000005\n",
    "}\n",
    "\n",
    "# The cost per token for each model output.\n",
    "MODEL_COST_PER_OUTPUT = {\n",
    "    \"gpt-3.5-turbo-0125\": 0.0000015,\n",
    "    \"gpt-4-turbo-2024-04-09\": 0.00003,\n",
    "    \"gpt-4o-2024-05-13\": 0.000015\n",
    "}\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-07-19T18:14:43.897176Z",
     "start_time": "2024-07-19T18:14:43.861428Z"
    }
   },
   "id": "60f0e21cfbcb050f",
   "execution_count": 26
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "  0%|          | 0/32 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f86f41c1fc9c42b289851756f32eb5f4"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  0.129525\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 1\n",
      "Total cost:  0.20917\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  0.334125\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 2\n",
      "Total cost:  0.411325\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  0.75716\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 3\n",
      "Total cost:  1.322475\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 4\n",
      "Total cost:  1.36727\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  1.672445\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  1.95275\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 5\n",
      "Total cost:  2.0527\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 6\n",
      "Total cost:  2.1470900000000004\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 7\n",
      "Total cost:  2.7226900000000005\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 8\n",
      "Total cost:  3.1455850000000005\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 9\n",
      "Total cost:  3.610175000000001\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  3.968925000000001\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 10\n",
      "Total cost:  4.298745000000001\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  4.486785000000001\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  4.5732300000000015\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  4.706190000000001\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  5.054345000000001\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 11\n",
      "Total cost:  5.2791700000000015\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  5.455825000000002\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 12\n",
      "Total cost:  5.995885000000001\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 13\n",
      "Total cost:  6.096265000000002\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  6.233800000000001\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 14\n",
      "Total cost:  6.7613400000000015\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  7.063930000000002\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  7.428230000000002\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  7.810855000000002\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  7.9558300000000015\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "Total cost:  8.058055000000001\n",
      "Upload a file with an assistants purpose\n",
      "start a messages\n",
      "start client running\n",
      "read  messages\n",
      "error 15\n",
      "Total cost:  8.251970000000002\n"
     ]
    }
   ],
   "source": [
    "error_cout = 0\n",
    "total_cost = 0\n",
    "save_path = \"./output_model/\"\n",
    "start = time.time()\n",
    "model = 'gpt-3.5-turbo-0125'\n",
    "# model = 'gpt-4o-2024-05-13'\n",
    "for line in tqdm(data[42:]):\n",
    "  name = line['name']\n",
    "  # print(name)\n",
    "  start = time.time()\n",
    "  with open(f\"./data/task/{name}.txt\", \"r\") as f:\n",
    "    description = f.read()\n",
    "  assistant = client.beta.assistants.create(\n",
    "    instructions=\"You are a data scientist. I have a data modeling task. You must give me the predicted results as a CSV file as detailed in the following content.  Don't ask me any questions. I provide you with three files. One is training data, one is test data. There is also a sample file for submission\",\n",
    "    model=model,\n",
    "    tools=[{\"type\": \"code_interpreter\"}],\n",
    "    # tool_resources={\n",
    "    #   \"code_interpreter\": {\n",
    "    #     \"file_ids\": [train_file.id, test_file.id, sample_file.id]\n",
    "    #   }\n",
    "    # }\n",
    "  )\n",
    "  files = os.listdir(f\"./data_resplit/{name}\")\n",
    "  print(\"Upload a file with an assistants purpose\")\n",
    "  train_file = client.files.create(\n",
    "    file=open(f\"./data_resplit/{name}/train.csv\", \"rb\"),\n",
    "    purpose='assistants'\n",
    "  )\n",
    "  \n",
    "  # Upload a file with an \"assistants\" purpose\n",
    "  test_file = client.files.create(\n",
    "    file=open(f\"./data_resplit/{name}/test.csv\", \"rb\"),\n",
    "    purpose='assistants'\n",
    "  )\n",
    "  files.remove('test.csv')\n",
    "  files.remove('train.csv')\n",
    "  sample_file = client.files.create(\n",
    "    file=open(f\"./data_resplit/{name}/{files[0]}\", \"rb\"),\n",
    "    purpose='assistants'\n",
    "  )\n",
    "  print(\"start a messages\")\n",
    "  thread = client.beta.threads.create(\n",
    "    messages=[\n",
    "      {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": description,\n",
    "        \"attachments\": [\n",
    "          {\n",
    "            \"file_id\": train_file.id,\n",
    "            \"tools\": [{\"type\": \"code_interpreter\"}]\n",
    "          },\n",
    "          {\n",
    "            \"file_id\": test_file.id,\n",
    "            \"tools\": [{\"type\": \"code_interpreter\"}]\n",
    "          },\n",
    "          {\n",
    "            \"file_id\": sample_file.id,\n",
    "            \"tools\": [{\"type\": \"code_interpreter\"}]\n",
    "          }\n",
    "        ]\n",
    "      }\n",
    "    ]\n",
    "  )\n",
    "  print(\"start client running\")\n",
    "  run = client.beta.threads.runs.create_and_poll(\n",
    "    thread_id=thread.id,\n",
    "    assistant_id=assistant.id,\n",
    "  )\n",
    "  prompt_tokens = run.usage.prompt_tokens \n",
    "  completion_tokens = run.usage.completion_tokens\n",
    "  cost = 0.03 + run.usage.completion_tokens * MODEL_COST_PER_OUTPUT[model] + run.usage.prompt_tokens * MODEL_COST_PER_INPUT[model]\n",
    "\n",
    "  print(\"read  messages\")\n",
    "  # cost = 0\n",
    "  if not os.path.exists(f\"{save_path}{model}\"):\n",
    "      os.makedirs(f\"{save_path}{model}\")\n",
    "  messages = \"\"\n",
    "  try:\n",
    "    messages = client.beta.threads.messages.list(\n",
    "      thread_id=thread.id\n",
    "    )\n",
    "    # print(messages)\n",
    "    submission_id = messages.data[0].content[0].text.annotations[0].file_path.file_id\n",
    "  \n",
    "    with open(f\"{save_path}{model}/{name}.csv\", \"wb\") as f:\n",
    "      f.write(client.files.content(submission_id).read())\n",
    "    client.files.delete(submission_id)\n",
    "  except:\n",
    "    error_cout += 1\n",
    "    print(f\"error {error_cout}\")\n",
    "  total_cost += cost\n",
    "  with open(f\"{save_path}{model}/{name}.json\", \"w\") as f:\n",
    "    json.dump({\"name\": name, \"model\": model, \"input\": prompt_tokens,\n",
    "                            \"output\": completion_tokens, \"cost\": cost, \"time\": time.time()-start, 'response': str(messages)}, f)\n",
    "  with open(f\"{save_path}{model}/{name}_message.txt\", \"w\") as f:\n",
    "    f.write(str(messages))\n",
    "  print(\"Total cost: \", total_cost)\n",
    "  client.files.delete(sample_file.id)\n",
    "  client.files.delete(train_file.id)\n",
    "  client.files.delete(test_file.id)\n",
    "  \n",
    "  client.beta.threads.delete(thread.id)\n",
    "  client.beta.assistants.delete(assistant.id)\n",
    "end = time.time()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-07-20T06:16:26.078235Z",
     "start_time": "2024-07-20T05:00:20.849193Z"
    }
   },
   "id": "fabb474a4383a861",
   "execution_count": 33
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
